import time
import os
import mujoco.viewer
import mujoco
import numpy as np
import torch
import yaml
import random
import math

MICRO_UR_ROOT_DIR = os.getcwd()

def euler_to_quaternion(roll, pitch, yaw):
    """
    将欧拉角 (Roll, Pitch, Yaw) 转换为四元数 (w, x, y, z)
    旋转顺序:ZYX (Yaw → Pitch → Roll)
    """
    output = np.zeros(4, dtype=np.float32)

    # 计算各轴旋转的半角三角函数
    cy = math.cos(yaw * 0.5)
    sy = math.sin(yaw * 0.5)
    cp = math.cos(pitch * 0.5)
    sp = math.sin(pitch * 0.5)
    cr = math.cos(roll * 0.5)
    sr = math.sin(roll * 0.5)

    # 计算四元数分量
    w = cr * cp * cy + sr * sp * sy
    x = sr * cp * cy - cr * sp * sy
    y = cr * sp * cy + sr * cp * sy
    z = cr * cp * sy - sr * sp * cy

    output[0] = w
    output[1] = x
    output[2] = y
    output[3] = z

    return output

def pd_control(target_q, q, kp, target_dq, dq, kd):
    """Calculates torques from position commands"""
    return (target_q - q) * kp + (target_dq - dq) * kd

if __name__ == "__main__":
    # get config file name from command line
    import argparse

    # 通过启动文件时 传入yaml文件名 来传递一些参数设置
    parser = argparse.ArgumentParser()
    parser.add_argument("config_file", type=str, help="config file name in the config folder")
    args = parser.parse_args()
    config_file = args.config_file

    with open(f"{MICRO_UR_ROOT_DIR}/src/fur_sim2sim/deploy_python/configs/{config_file}", "r") as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
        policy_path = config["policy_path"].replace("{MICRO_UR_ROOT_DIR}", MICRO_UR_ROOT_DIR)
        xml_path = config["xml_path"].replace("{MICRO_UR_ROOT_DIR}", MICRO_UR_ROOT_DIR)

        simulation_duration = config["simulation_duration"]
        simulation_dt = config["simulation_dt"]
        control_decimation = config["control_decimation"]

        kps = np.array(config["kps"], dtype=np.float32)
        kds = np.array(config["kds"], dtype=np.float32)

        stand_kps = np.array(config["stand_kps"], dtype=np.float32)
        stand_kds = np.array(config["stand_kds"], dtype=np.float32)

        default_angles = np.array(config["default_angles"], dtype=np.float32)
        initial_angles = np.array(config["initial_angles"], dtype=np.float32)

        dof_pos_scale = config["dof_pos_scale"]
        dof_vel_scale = config["dof_vel_scale"]
        action_scale = config["action_scale"]
        cmd_scale = np.array(config["cmd_scale"], dtype=np.float32) # 对于一些观测值的缩放

        num_actions = config["num_actions"]
        num_obs = config["num_obs"]
        
    # define context variables
    action = np.zeros(num_actions, dtype=np.float32)
    target_dof_pos = np.zeros(num_actions, dtype=np.float32)
    obs = np.zeros(num_obs, dtype=np.float32)

    counter = 0

    # Load robot model
    m = mujoco.MjModel.from_xml_path(xml_path) 
    d = mujoco.MjData(m)
    m.opt.timestep = simulation_dt

    # load policy
    policy = torch.jit.load(policy_path)

    percent = 0.0
    stand_duration = 5.0
    pose_cmd = np.zeros(6, dtype=np.float32)
    stand_flag = False

    # 在MuJoCo中运行一个物理仿真，并通过被动查看器实时显示
    with mujoco.viewer.launch_passive(m, d) as viewer:
        # Close the viewer automatically after simulation_duration wall-seconds.
        start = time.time()
        while viewer.is_running() and time.time() - start < simulation_duration:
            step_start = time.time()

            if percent < 1:
                if stand_flag ==  False:
                    now_dof_pos = d.qpos[:]
                    stand_flag = True
                for i in range(6):
                    target_dof_pos[i] = (1 - percent)*(now_dof_pos[i]) + percent*default_angles[i]
                tau = pd_control(target_dof_pos, d.qpos[:], stand_kps, np.zeros_like(stand_kds), d.qvel[:], stand_kds)
                # 设置控制力矩
                d.ctrl[:] = tau
                percent += 1 / (stand_duration * (1 / simulation_dt))
                

            else:
                # 随机pose命令
                if counter % 2500 == 0: # 每5s一个新命令 random.uniform
                    pose_cmd[0] = random.uniform(0.15, 0.30)
                    pose_cmd[1] = random.uniform(-0.20, 0.20)
                    pose_cmd[2] = random.uniform(0.20, 0.40)
                    pose_cmd[3] = random.uniform(-1.57, 1.57)
                    pose_cmd[4] = random.uniform(-1.57, 1.57)
                    pose_cmd[5] = random.uniform(-1.57, 1.57)

                if counter % control_decimation == 0:
                    # create observation
                    qpos = d.qpos[:]
                    qvel = d.qvel[:]
                    qpos = (qpos - default_angles) * dof_pos_scale
                    qvel = qvel * dof_vel_scale

                    pose_quat_cmd = np.zeros(7, dtype=np.float32)
                    pose_quat_cmd[0] = pose_cmd[0]
                    pose_quat_cmd[1] = pose_cmd[1]
                    pose_quat_cmd[2] = pose_cmd[2]
                    pose_quat_cmd[3:] = euler_to_quaternion(pose_cmd[3], pose_cmd[4], pose_cmd[5])

                    obs[:6] = qpos
                    obs[6:12] = qvel
                    obs[12:19] = pose_quat_cmd * cmd_scale
                    obs[19 : 19 + num_actions] = action
                    obs_tensor = torch.from_numpy(obs).unsqueeze(0) # 拼接观测向量
                    obs_tensor = torch.clip(obs_tensor, min=-100.0, max=100.0)
                    #print(obs_tensor)
                    # policy inference
                    action = policy(obs_tensor).detach().numpy().squeeze()
                    #print(action)
                    # transform action to target_dof_pos
                    target_dof_pos = action * action_scale + default_angles

                counter += 1

                tau = pd_control(target_dof_pos, d.qpos[:], kps, np.zeros_like(kds), d.qvel[:], kds)
                # 设置控制力矩
                d.ctrl[:] = tau

            mujoco.mj_step(m, d)

            # Pick up changes to the physics state, apply perturbations, update options from GUI.
            viewer.sync()

            # Rudimentary time keeping, will drift relative to wall clock.
            # 计算当前步进剩余时间 确保仿真步进时间与物理时间匹配
            time_until_next_step = m.opt.timestep - (time.time() - step_start)
            if time_until_next_step > 0:
                time.sleep(time_until_next_step)

            




